# Copyright ClusterHQ Inc.  See LICENSE file for details.

"""
Testing utilities for ``flocker.ca``.
"""

import atexit
from tempfile import mkdtemp
from shutil import rmtree

from OpenSSL.crypto import X509Extension

from twisted.python.filepath import FilePath

from pyrsistent import PClass, field

from ..ca import (
    RootCredential, ControlCredential, NodeCredential, UserCredential,
)


def assert_has_extension(test, credential, name, value):
    """
    Assert that the ``X509Extension`` with the matching name from the
    certificate has the given value.

    :param TestCase test: The current test.
    :param FlockerCredential certificate: Credential whose certificate we
        should inspect.
    :param bytes name: The name of the extension.
    :param bytes value: The data encoded in the extension.

    :raises AssertionError: If the extension is not found or has the wrong
        value.
    """
    expected = X509Extension(name, False, value)
    x509 = credential.certificate.original
    values = []
    for i in range(x509.get_extension_count()):
        extension = x509.get_extension(i)
        if extension.get_short_name() == name:
            values.append(extension.get_data())
    test.assertIn(expected.get_data(), values)


class CredentialSet(PClass):
    """
    A full set of credentials for a CA.

    :ivar FilePath path: Directory where credentials reside.
    :ivar RootCredential root: The CA root credential.
    :ivar ControlCredential control: A control service credential using an
        IP for hostname.
    :ivar ControlCredential control_dns: A control service credential using a
        DNS name for hostname.
    :ivar UserCredential user: A user credential.
    :ivar NodeCredential node: A CA root credentials.
    """
    path = field()
    root = field()
    control = field()
    control_dns = field()
    user = field()
    node = field()

    @staticmethod
    def create():
        """
        :return: A new ``CredentialSet``.
        """
        directory = FilePath(mkdtemp())
        atexit.register(rmtree, directory.path)
        root = RootCredential.initialize(directory, b"mycluster")
        user = UserCredential.initialize(directory, root, u"allison")
        node = NodeCredential.initialize(directory, root)
        control = ControlCredential.initialize(directory, root, b"127.0.0.1")
        control_dns = ControlCredential.initialize(
            directory, root, b"localhost")
        return CredentialSet(path=directory,
                             root=root, user=user, node=node, control=control,
                             control_dns=control_dns)

    def copy_to(self, directory, control=False, user=False, node=False):
        """
        Copy certificates to destination directory with naming convention used
        by certificate-loading code (control service, node agents and
        clients).

        Always copies CA certificate.

        :param FilePath directory: Destination for files.
        :param bool control: Whether to copy control service certificate/key.
        :param bool node: Whether to copy node certificate/key.
        :param bool user: Whether to copy API client certificate/key.
        """
        self.path.child(b"cluster.crt").copyTo(directory.child(b"cluster.crt"))
        if control:
            self.path.child(b"control-127.0.0.1.crt").copyTo(
                directory.child(b"control-service.crt"))
            self.path.child(b"control-127.0.0.1.key").copyTo(
                directory.child(b"control-service.key"))
        if user:
            self.path.child(b"allison.crt").copyTo(
                directory.child(b"user.crt"))
            self.path.child(b"allison.key").copyTo(
                directory.child(b"user.key"))
        if node:
            uuid = self.node.uuid
            self.path.child(b"{}.crt".format(uuid)).copyTo(
                directory.child(b"node.crt"))
            self.path.child(b"{}.key".format(uuid)).copyTo(
                directory.child(b"node.key"))


# Cached CredentialSets generated by get_sets():
_credential_set1 = None
_credential_set2 = None


def get_credential_sets():
    """
    Get a pair of cached sets of CA-specific credentials.

    Since they are cached across tests they should not be modified.

    :return: Tuple, a pair of ``CredentialSet`` instances.
    """
    # Yes, we shouldn't use global in general, but this is test code and
    # it's the easiest way to have cached value.
    global _credential_set2  # pylint: disable=global-statement
    global _credential_set1  # pylint: disable=global-statement
    if _credential_set1 is None:
        _credential_set1 = CredentialSet.create()
        _credential_set2 = CredentialSet.create()
    return _credential_set1, _credential_set2
